# -*- coding: utf-8 -*-
"""
-------------------------------------------------
File Name： bbox_transform
Description :
Author : 'li'
date： 2022/1/28
-------------------------------------------------
Change Activity:
2022/1/28:
-------------------------------------------------
"""
import numpy as np
import torch
import cv2


def obb2poly(obboxes):
    """
    Trans rbox format to poly format.
    Args:
        obboxes (array/tensor): (num_gts, [cx cy l s θ]) θ∈[-pi/2, pi/2)

    Returns:
        polys (array/tensor): (num_gts, [x1 y1 x2 y2 x3 y3 x4 y4])
    """
    if isinstance(obboxes, torch.Tensor):
        center, w, h, theta = obboxes[:, :2], obboxes[:, 2:3], obboxes[:, 3:4], obboxes[:, 4:5]
        cos, sin = torch.cos(theta), torch.sin(theta)

        vector1 = torch.cat(
            (w / 2 * cos, -w / 2 * sin), dim=-1)
        vector2 = torch.cat(
            (-h / 2 * sin, -h / 2 * cos), dim=-1)
        point1 = center + vector1 + vector2
        point2 = center + vector1 - vector2
        point3 = center - vector1 - vector2
        point4 = center - vector1 + vector2
        order = obboxes.shape[:-1]
        return torch.cat((point1, point2, point3, point4), dim=-1).reshape(*order, 8)
    else:
        center, w, h, theta = np.split(obboxes, (2, 3, 4), axis=-1)
        cos, sin = np.cos(theta), np.sin(theta)

        vector1 = np.concatenate(
            [w / 2 * cos, -w / 2 * sin], axis=-1)
        vector2 = np.concatenate(
            [-h / 2 * sin, -h / 2 * cos], axis=-1)

        point1 = center + vector1 + vector2
        point2 = center + vector1 - vector2
        point3 = center - vector1 - vector2
        point4 = center - vector1 + vector2
        order = obboxes.shape[:-1]
        return np.concatenate(
            [point1, point2, point3, point4], axis=-1).reshape(*order, 8)


def np_obb2poly(cxcy, wh, delta):
    return cv2.boxPoints((cxcy, wh, delta))


def degree2radian(degree):
    """
    180 2 pi
    @param degree:
    @return:
    """
    return np.deg2rad(degree)


def radian2degree(rad):
    """

    @param rad:
    @return:
    """
    return np.rad2deg(rad)


def poly2obb(poly):
    res = cv2.minAreaRect(poly.astype(np.int32))
    return res


def poly2hbb(poly):
    """

    Args:
        poly:

    Returns:

    """
    ws, hs = poly[..., 0], poly[..., 1]
    lt = [ws.min(), hs.min()]  # left top point
    rd = [ws.max(), hs.max()]
    return np.array([lt, rd])


def get_poly_height_and_width(poly):
    """

    Args:
        poly:

    Returns:

    """
    ws, hs = poly[..., 0], poly[..., 1]
    lt = [ws.min(), hs.min()]  # left top point
    rd = [ws.max(), hs.max()]
    width = rd[0] - lt[0]
    height = rd[1] - lt[1]
    return width, height


def boxes_cxcywh2xyxy(x):
    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
    y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left x
    y[:, 1] = x[:, 1] - x[:, 3] / 2  # top left y
    y[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right x
    y[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right y
    return y


def boxes_xyxy2cxcywh(bboxes):
    """

    Args:
        bboxes:

    Returns:

    """
    bboxes[:, 2] = bboxes[:, 2] - bboxes[:, 0]
    bboxes[:, 3] = bboxes[:, 3] - bboxes[:, 1]
    return bboxes


def box_cxcywh_to_xyxy(x):
    """

    Args:
        x:

    Returns:

    """
    x_c, y_c, w, h = x.unbind(-1)
    b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
         (x_c + 0.5 * w), (y_c + 0.5 * h)]
    return torch.stack(b, dim=-1)


def box_xyxy2cxcywh(box):
    """

    Args:
        box:

    Returns:

    """
    x0, y0, x1, y1 = box.unbind(-1)
    b = [(x0 + x1) / 2, (y0 + y1) / 2,
         (x1 - x0), (y1 - y0)]
    return torch.stack(b, dim=-1)


def cxcywh2ploy(cxcywh):
    xyxy = boxes_cxcywh2xyxy(cxcywh)
    return xyxy2ploy(xyxy)


def single_xyxy2poly(xyxy):
    x1, y1, x2, y2 = xyxy
    p = [[x1, y1], [x2, y1], [x2, y2], [x1, y2]]
    p = np.array(p)
    return p


def ploy2xyxy(ploy):
    """

    Args:
        ploy:(point_size,2)  (w,h)

    Returns:

    """
    _ploy = np.array(ploy)
    w_min, h_min = _ploy[:,0].min(), _ploy[:,1].min()
    w_max, h_max = _ploy[:,0].max(), _ploy[:,1].max()
    return [[w_min, h_min], [w_max, h_max]]


def xyxy2ploy(xyxy_lst):
    """

    Args:
        xyxy_lst:

    Returns:

    """
    ps = []
    for xyxy in xyxy_lst:
        p = single_xyxy2poly(xyxy)
        ps.append(p)
    return np.stack(ps, 0)


def enlarge_bbox(enlarge_width, enlarge_height, bbox, shape):
    """

    Args:
        enlarge_width:
        enlarge_height:
        bbox:
        shape:

    Returns:

    """
    if len(bbox.shape) == 1:
        raise Exception('todo')
    p1, p2, p3, p4 = bbox[0], bbox[1], bbox[2], bbox[3]
    p1 = [p1[0] - enlarge_width, p1[1] - enlarge_height]
    p2 = [p2[0] + enlarge_width, p2[1] - enlarge_height]
    p3 = [p3[0] + enlarge_width, p3[1] + enlarge_height]
    p4 = [p4[0] - enlarge_width, p4[1] + enlarge_height]
    enlarged_bbox = [p1, p2, p3, p4]
    img_height, img_width = shape[0], shape[1]
    for point in enlarged_bbox:
        width, height = point
        if width < 0:
            point[0] = 0
        if height < 0:
            point[1] = 0
        if height > (img_height - 1):
            point[1] = img_height - 1
        if width > (img_width - 1):
            point[0] = (img_width - 1)
    enlarged_bbox = np.array(enlarged_bbox)
    return enlarged_bbox
